package leetcode

import (
	"math/rand"
	"time"
)

// https://leetcode.com/problems/random-pick-with-weight/
type Solution struct {
	preSum []int
	seed   *rand.Rand
}

func Constructor(w []int) Solution {
	seed := rand.New(rand.NewSource(time.Now().UnixNano()))
	preSum := make([]int, len(w)+1)
	for i := 1; i < len(preSum); i++ {
		preSum[i] = preSum[i-1] + w[i-1]
	}
	return Solution{preSum: preSum, seed: seed}
}

func (t *Solution) PickIndex() int {
	target := t.seed.Intn(t.preSum[len(t.preSum)-1]) + 1
	return t.LeftBound(target) - 1
}

// return the index of element which first greater or equal target.
func (t *Solution) LeftBound(target int) int {
	left, right := 0, len(t.preSum)-1
	for mid := (left + right) / 2; left <= right; mid = (left + right) / 2 {
		if t.preSum[mid] < target {
			left = mid + 1
		} else {
			if mid > 0 && t.preSum[mid-1] < target {
				return mid
			}
			right = mid - 1
		}
	}
	return -1
}
